"""
生成示例数据
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from datetime import datetime, timedelta


def generate_random_prices(start_date, periods=252, start_price=100, volatility=0.01, trend=0.0001):
    """
    生成随机价格数据
    
    Parameters
    ----------
    start_date : datetime
        起始日期
    periods : int, optional
        数据点数量，默认为252（一年的交易日）
    start_price : float, optional
        起始价格，默认为100
    volatility : float, optional
        波动率，默认为0.01
    trend : float, optional
        趋势，默认为0.0001
        
    Returns
    -------
    pandas.DataFrame
        包含价格数据的DataFrame
    """
    # 生成日期范围
    dates = [start_date + timedelta(days=i) for i in range(periods)]
    
    # 生成随机价格
    returns = np.random.normal(trend, volatility, periods)
    price_index = (1 + returns).cumprod()
    prices = start_price * price_index
    
    # 生成OHLC数据
    opens = prices * (1 + np.random.normal(0, volatility/2, periods))
    highs = np.maximum(prices, opens) * (1 + np.abs(np.random.normal(0, volatility, periods)))
    lows = np.minimum(prices, opens) * (1 - np.abs(np.random.normal(0, volatility, periods)))
    closes = prices
    volumes = np.random.lognormal(mean=10, sigma=1, size=periods).astype(int)
    
    # 创建DataFrame
    df = pd.DataFrame({
        'datetime': dates,
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes,
        'volume': volumes
    })
    
    return df


def generate_sine_wave_prices(start_date, periods=252, start_price=100, amplitude=5, frequency=0.05):
    """
    生成正弦波价格数据
    
    Parameters
    ----------
    start_date : datetime
        起始日期
    periods : int, optional
        数据点数量，默认为252（一年的交易日）
    start_price : float, optional
        起始价格，默认为100
    amplitude : float, optional
        振幅，默认为5
    frequency : float, optional
        频率，默认为0.05
        
    Returns
    -------
    pandas.DataFrame
        包含价格数据的DataFrame
    """
    # 生成日期范围
    dates = [start_date + timedelta(days=i) for i in range(periods)]
    
    # 生成正弦波价格
    x = np.arange(periods)
    trend = 0.05 * x  # 线性趋势
    noise = np.random.normal(0, 1, periods)  # 随机噪声
    closes = start_price + trend + amplitude * np.sin(frequency * x) + noise
    
    # 生成OHLC数据
    volatility = 0.01
    opens = closes * (1 + np.random.normal(0, volatility/2, periods))
    highs = np.maximum(closes, opens) * (1 + np.abs(np.random.normal(0, volatility, periods)))
    lows = np.minimum(closes, opens) * (1 - np.abs(np.random.normal(0, volatility, periods)))
    volumes = np.random.lognormal(mean=10, sigma=1, size=periods).astype(int)
    
    # 创建DataFrame
    df = pd.DataFrame({
        'datetime': dates,
        'open': opens,
        'high': highs,
        'low': lows,
        'close': closes,
        'volume': volumes
    })
    
    return df


def save_data(df, filename):
    """
    保存数据到CSV文件
    
    Parameters
    ----------
    df : pandas.DataFrame
        数据
    filename : str
        文件名
    """
    df.to_csv(filename, index=False)
    print(f"数据已保存到 {filename}")


def plot_data(df, title="价格数据"):
    """
    绘制价格数据
    
    Parameters
    ----------
    df : pandas.DataFrame
        数据
    title : str, optional
        图表标题，默认为"价格数据"
    """
    plt.figure(figsize=(12, 6))
    plt.plot(df['datetime'], df['close'], label='收盘价')
    plt.plot(df['datetime'], df['open'], label='开盘价', alpha=0.5)
    plt.plot(df['datetime'], df['high'], label='最高价', alpha=0.3)
    plt.plot(df['datetime'], df['low'], label='最低价', alpha=0.3)
    plt.title(title)
    plt.xlabel('日期')
    plt.ylabel('价格')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='生成示例数据')
    parser.add_argument('--type', type=str, choices=['random', 'sine'], default='random', help='数据类型: random(随机) 或 sine(正弦波)')
    parser.add_argument('--periods', type=int, default=252, help='数据点数量')
    parser.add_argument('--start_price', type=float, default=100, help='起始价格')
    parser.add_argument('--output', type=str, default='sample_data.csv', help='输出文件名')
    parser.add_argument('--plot', action='store_true', help='是否绘制图表')
    
    args = parser.parse_args()
    
    start_date = datetime.now() - timedelta(days=args.periods)
    
    if args.type == 'random':
        df = generate_random_prices(start_date, args.periods, args.start_price)
    else:
        df = generate_sine_wave_prices(start_date, args.periods, args.start_price)
    
    save_data(df, args.output)
    
    if args.plot:
        plot_data(df, f"{args.type.capitalize()} 价格数据") 